import copy
import math

from DrawingAPIs.bar_drawer import *

def ParallelDegree(pDeg_file):
    streams = {}
    with open(pDeg_file, newline='') as csvfile:
        rows = csv.reader(csvfile, delimiter=',')
        for row in rows:
            streams[row[0].split('FW', 1)[0][:-1]] = int(row[1])

    return streams

def DrawStreamFig(streams_data, conv_layers, devices, out_dir):
    # Data format: {device: {model name: {layer name: #streams}}}

    # Data preprocessing.
    for net in streams_data[devices[0]].keys():
        x_ticks = np.arange(2, 2 * (len(conv_layers[net]) + 1), 2)
        x_ticklabels = []
        bar_w = 0.45
        y_maxval = 0
        stream_y_values = []
        for dev in devices:
            temp_conv_streams = []
            for layer in conv_layers[net]:
                temp_conv_streams.append(streams_data[dev][net][layer])

            # Add device-specific data to the y_value_set
            y_maxval = max(y_maxval, max(temp_conv_streams))
            stream_y_values.append(temp_conv_streams)
        y_maxval += 1
        for i in range(0, len(conv_layers[net])):
            x_ticklabels.append('conv' + str(i + 1))

        # print("Drawing the stream result of the {0} model...".format(net))
        DrawBarFig(x_ticks, x_ticklabels, stream_y_values, devices, '', '#streams', bar_w,
                   x_min=0.8, x_max=(x_ticks[-1] + 1), y_min=0, y_max=y_maxval,
                   folder=(out_dir + '/Streams'), filename=('exp_optStreams_' + net), allow_legend=True)

def DrawStreamFigs(args, devices, net_layer_used):
    # Work folder: args.root[0] + <device>/IntraLayer/<network>/ParallelDegree/ForwardPDeg_rec.csv
    streams_data = {}
    for dev in devices:
        if dev not in streams_data.keys():
            streams_data.setdefault(dev, {})

        dev_folder = args.root[0] + '/' + dev + '/IntraLayer'
        nets = []
        for net in os.listdir(dev_folder):
            if net.lower() == "googlenet":
                continue
            
            if net not in nets:
                nets.append(net)

            if net not in streams_data[dev].keys():
                streams_data[dev].setdefault(net, {})

            pdegree_file = dev_folder + '/' + net + '/ParallelDegree/ForwardPDeg_rec.csv'
            streams_data[dev][net] = ParallelDegree(pdegree_file)

    # Draw stream-related figures.
    DrawStreamFig(streams_data, net_layer_used, devices, args.output[0])
